#include <iostream>
#include <fstream>
#include <sys/mman.h>
#include <string.h>
#include <gtest/gtest.h>
#include <vector>
#include "Common.h"
#include "Segment.h"

#if defined(USE_NPU)
#include "acl/acl.h"
#endif

#ifdef USE_GPU
#include "cuda_runtime.h"
#endif

using namespace std;

#define MEM_SIZE (1024 * 1024 * 1024)

class TestSegment : public ::testing::Test {
protected:
  void SetUp() override {
  }

  void TearDown() override {
  }

};

void prepareMemAddr(char **extMemAddr, char **segMemAddr, void **deviceMemAddr)
{
    *extMemAddr = (char *)malloc(MEM_SIZE);
    assert(nullptr != *extMemAddr);
    *segMemAddr = (char *)malloc(MEM_SIZE);
    assert(nullptr != *segMemAddr);
#if defined(USE_NPU)
    assert(0 == aclrtSetDevice(0));
    assert(0 == aclrtMalloc(deviceMemAddr, MEM_SIZE, ACL_MEM_MALLOC_NORMAL_ONLY));
    assert(nullptr != *deviceMemAddr);
#elif defined(USE_GPU)
    assert(0 == cudaMalloc(deviceMemAddr, MEM_SIZE));
    assert(nullptr != *deviceMemAddr);
#else
    *deviceMemAddr = (char *)malloc(MEM_SIZE);
    assert(nullptr != *deviceMemAddr);
#endif
}

void freeMemAddr(char *extMemAddr, char *segMemAddr, void *deviceMemAddr)
{
    free(extMemAddr);
    free(segMemAddr);
#if defined(USE_NPU)
    aclrtFree(deviceMemAddr);
#elif defined(USE_GPU)
    cudaFree(deviceMemAddr);
#else
    free(deviceMemAddr);
#endif
}

void checkSharedMemory(char *mapAddr, const uint64_t off, const uint64_t len, char tgt) {
    for (uint64_t i = 0; i < len; ++i) {
        if (mapAddr[i + off] != tgt) {
            std::cout << "first mismatch idx(" << i + off << "):";
            assert(tgt == mapAddr[i]);
            return;
        }
    }
}

TEST(TestSegment, Normal) {
    char *extMemAddr = nullptr;
    char *segMemAddr = nullptr;
    void *deviceMemAddr = nullptr;
    prepareMemAddr(&extMemAddr, &segMemAddr, &deviceMemAddr);

    Segment segment;
    segment.init(0, 1, 14688256, (uint8_t *)segMemAddr, 7340032, 10, true, false, "", false);

    DeviceType type;
#if defined(USE_NPU)
    type = DEVICE_TYPE_NPU;
#elif defined(USE_GPU)
    type = DEVICE_TYPE_GPU;
#else
    type = DEVICE_TYPE_CPU;
#endif

    std::fill(extMemAddr, extMemAddr + MEM_SIZE, 'G');
    assert(KVC_OK == segment.writeIn(type, extMemAddr, 0, 7340032));
    assert(KVC_OK == segment.copyTo(type, deviceMemAddr, 7077888, 131072));
    std::fill(segMemAddr, segMemAddr + MEM_SIZE, '0');
    assert(KVC_OK == segment.writeIn(type, deviceMemAddr, 7077888, 131072));
    checkSharedMemory(segMemAddr, 7077888, 131072, 'G');
}

int main(int argc, char **argv) {
    testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}
